import wx
import numpy as np
import matplotlib as mpl
import wx.lib.agw.customtreectrl as CT
from matplotlib.figure import Figure
from PlotPanel2 import PlotPanel, My_Axes
import ConfigParser
import state
from dissociation import DissociationLimits
import tracer as tr
import argparse

mpl.projections.register_projection(My_Axes)

class TreePlotPanel(wx.Panel):
    def __init__(self, parent, filename):
        wx.Panel.__init__(self, parent)
        self.configfile = filename

        self.window = wx.SplitterWindow(self, wx.ID_ANY, style=wx.SP_3D | wx.SP_BORDER)
        self.read_config_file(filename)
        #Setup basic plot objects
        self.Fig = Figure()
        self.ax = self.Fig.add_subplot(111, projection="My_Axes")
        self.Title = self.name

        self.initPlot()
        self.treepanel = wx.Panel(self.window, wx.ID_ANY)
        self.plotpanel = PecPanel(self.window, self.Fig, self.winvals, self.states, self.dislimits, self.units)

        self.tree = CT.CustomTreeCtrl(self.treepanel,
                                      agwStyle=wx.TR_DEFAULT_STYLE | CT.TR_AUTO_CHECK_CHILD)
        self.root = self.tree.AddRoot("States", ct_type = 1)
        self.tree.Expand(self.root)
        self.tree.CheckItem2(self.root)
        self.treeitems = []
        for item in self.states:
            self.addTreeItem(item)
        #Bindings
        self.tree.Bind(CT.EVT_TREE_ITEM_CHECKED, self.OnSelect)

        self.makeLayout()
        
    def read_config_file(self, filename):
        '''Gets config file information'''
        self.units = {}

        self.datafiles = {}
        self.styles = {}
        self.axprops = {} #Window Properties
        self.vprops = {}
        self.fstring = "{}.txt"

        parser = ConfigParser.ConfigParser()
        parser.optionxform = str
        parser.read(filename)

        self.name = parser.get("Molecule", "Name")

        options = parser.options("Units")
        for key in options:
            self.units[key] = parser.get("Units", key)

        options = parser.options("Files")
        for key in options:
            if key == "Format":
                self.fstring = parser.get("Files", key)
            else:
                self.datafiles[key] = parser.get("Files", key).split(", ")

        options = parser.options("Styles")
        for key in options:
            self.styles[key] = parser.get("Styles", key)

        options = parser.options("Window")
        for key in options:
            self.axprops[key] = parser.getfloat("Window", key)

        options = parser.options("Virtual")
        for key in options:
            self.vprops[key] = parser.getfloat("Virtual", key)
        self.winvals = [ self.vprops['xmin'], self.vprops['xmax'],
                         self.vprops['ymin'], self.vprops['ymax'] ]
            
    def initPlot(self):
        '''Creates the starting figure''' 
        self.ax.set_xlim(self.axprops['xmin'], self.axprops['xmax'])
        self.ax.set_xticks( np.arange(self.vprops['xmin'], self.vprops['xmax'], self.axprops['xticks']) )
        self.ax.set_yticks( np.arange(self.vprops['ymin'],self.vprops['ymax'],self.axprops['yticks']) )
        self.ax.set_ylim(self.axprops['ymin'], self.axprops['ymax'])
        self.ax.set_win_props(self.winvals)
        self.generate_styles()
        self.states = []
        self.dislimits = DissociationLimits(self.name)
        self.dislimits.draw(self.ax, self.units)
        self.dislimits.update(self.ax)
        #Create a new State for each datafile and add to states
        for key in self.datafiles:
            self.addState(self.datafiles[key], key, self.csdict[key])
        
    
    def addState(self, fileinfo, label, style):
        '''Adds new states to plot and tree control'''
        #Create State and plot
        additem = state.State(fileinfo, label, style, self.name)
        additem.plotPEC(self.ax, self.units)
        self.states.append(additem)
        return additem

    def removeState(self, index):
        '''Deletes curves and annotations corresponding to a given state index,
         then deletes state'''

        for i, curve in enumerate(self.states[index].curves):
            curve.remove()
            if self.states[index].annotes[i]:
                self.states[index].annotes[i].set_visible(False)
                self.states[index].annotes[i].remove()
        self.states[index].curves = []
        self.states[index].annotes = []
        self.tree.Delete(self.treeitems.pop(index))
        return self.states.pop(index)    
    
    def addTreeItem(self, state):
        parent = self.tree.AppendItem(self.root, state.label,
                                         ct_type = 1, data = state)
        self.treeitems.append(parent)
        attr = parent.Attr()
        attr.SetTextColour(state.color)
        for i, curve in enumerate(state.curves):
            child = self.tree.AppendItem(parent, '('+ str(i+1) + ')',
                                         ct_type = 1, data = curve)
            cattr = child.Attr()
            cattr.SetTextColour(state.color)

        self.tree.CheckChilds(parent)
        self.tree.CheckItem2(parent)
  
    def generate_styles(self):
        """Creates dictionary of colors/styles/linewidths corresponding
        to different states"""
        self.csdict = {}
        for key in self.datafiles:
              self.csdict[key] = self.get_style(key)

    def get_style(self, key):
        '''Gets a style according to the settings in the config file'''
        legend = {'1': 'Singlets', '2': 'Doublets', '3': 'Triplets',
                  'u' : 'Ungerade', 'g': 'Gerade'}
        color = self.styles[key[1]]
        style = self.styles[legend[key[0]]]
        width = float(self.styles[legend[key[-1]]])
        if key.find('u') != -1:
            width = float(self.styles['Ungerade'])
        else:
            width = float(self.styles['Gerade'])
        return [color, style, width]

    def setStates(self, newfiles):
        '''Reset states list with data from newfiles.'''
     
        #Ignore states that don't need to be replotted
        toplot = list(set(newfiles) - set(self.datafiles))
        #Find which states need to be deleted
        todel = list(set(self.datafiles) - set(newfiles))

        #Delete States
        for k in todel:
            for i, state in enumerate(self.states):
                if state.label == k:
                    self.removeState(i)

        self.plotpanel.redraw()
        bg = self.plotpanel.canvas.copy_from_bbox(self.ax.bbox)
        self.plotpanel.canvas.restore_region(bg)

        #Add States
        for k in toplot:
            self.csdict[k] = self.get_style(k)
            newst = self.addState(newfiles[k], k, self.csdict[k])
            for curve in newst.curves:
                self.ax.draw_artist(curve)
        
        self.plotpanel.canvas.blit(self.ax.bbox)
        self.datafiles = newfiles
        
    def resetPlot(self, newfiles, newunits):
        self.units = newunits
        self.datafiles = newfiles
        self.plotpanel.clearPlot()
        self.initPlot()
        self.plotpanel.redraw()

    def resetUnits(self, newunits):
        self.units = newunits
        
        self.dislimits.draw(self.ax, self.units)
        self.dislimits.update(self.ax)
        
        for state in self.states:
            state.plotPEC(self.ax, self.units)
            

    def OnSelect(self, event):
        '''Event Handler for check/uncheck'''
        selection = event.GetItem()
        currentstate = selection.IsChecked()
        root = self.tree.GetRootItem()
        if selection == root:
            for state in self.states:
                state.WaxOnWaxOff(currentstate)
            self.tree.CheckChilds(selection, currentstate)
        else:
            parent = selection.GetParent()
            if parent == root:
                state = selection.GetData()
                state.WaxOnWaxOff(currentstate)
                self.tree.CheckChilds(selection, currentstate)
            else:
                state = parent.GetData()
                line = selection.GetData()
                flag = not line.get_visible()
                state.OnOff(line, flag)
        self.plotpanel.redraw()

    def makeLayout(self):

        self.vbox = wx.BoxSizer(wx.VERTICAL)
        pltbox = wx.BoxSizer(wx.VERTICAL)
        self.treebox = wx.BoxSizer(wx.VERTICAL)
        self.treebox.Add(self.tree, 1, wx.EXPAND, 0)
        self.treepanel.SetSizer(self.treebox)
        pltbox.Add(self.plotpanel, 1, wx.EXPAND, 0)
        #self.plotpanel.SetSizer(pltbox)
        self.window.SplitVertically(self.treepanel, self.plotpanel, 150)
        self.vbox.Add(self.window, 1, wx.EXPAND, 0)
        self.SetSizer(self.vbox)
        self.vbox.Fit(self)
        self.Layout()
        
    def write_config_file(self, filename):
        '''Writes Config file information'''
        parser = ConfigParser.ConfigParser()
        parser.optionxform = str

        parser.add_section("Molecule")
        parser.set("Molecule", "Name", self.name)

        parser.add_section("Units")
        for key in self.units:
            parser.set("Units", key, self.units[key])

        parser.add_section("Files")
        parser.set("Files", "Format", self.fstring)
        for key in self.datafiles:
            line = ''
            for item in self.datafiles[key][:-1]:
                line += item + ", "
            line += self.datafiles[key][-1]
            parser.set("Files", key, line)

        parser.add_section("Styles")
        for key in self.styles:
            parser.set("Styles", key, self.styles[key])

        parser.add_section("Window")
        for key in self.plotpanel.axprops:
            parser.set("Window", key, self.plotpanel.axprops[key])

        parser.add_section("Virtual")
        for key in self.plotpanel.vprops:
            parser.set("Virtual", key, self.plotpanel.vprops[key])

        with open(filename, 'wb') as configfile:
            parser.write(configfile)

class PecPanel(PlotPanel):
    def __init__(self, parent, Fig, vwin, states, dislimits, units, *args, **kwds):
        #read config file
        self.parent = parent
        self.states = states
        self.diss = dislimits
        self.units = units
        PlotPanel.__init__(self, parent, Fig, vwin, *args, **kwds)

        #Initialize animations
        self.tracer = None
        self.dragger = None
        self.labeller = None
        self.annoline, = self.ax.plot((0,0),(1,1), color = 'k')
        self.annoline.set_visible(False)

    def labelAxes(self):
        self.unitex = {'angstrom': '$\AA$', 'AU': 'A.U.', 'nm': 'nm', 'wavenumbers': 'cm$^{-1}$' }
        self.ax.set_xlabel("Radius (" + self.unitex[self.units['Radius']]+")")
        self.ax.set_ylabel("Energy (" + self.unitex[self.units['Energy']]+")")

    def OnMove(self, event):
        '''Mouse motion event handler'''

        if self.tracer:  #Allows user to trace currently selected curve
            if event.inaxes:
                self.tracer.move(event.xdata, event.ydata)
                update = "Internuclear Distance = " + self.tracer.text[0]  + ' ' \
                         + self.units['Radius'] + "  Energy = " + self.tracer.text[1] + ' '  \
                         + self.units['Energy']
                self.statusbar.SetStatusText(update)
        elif self.labeller: #Animates line for labeling feature
            self.labeller.move(event.xdata, event.ydata)
        elif self.dragger: #Allows annotations to be dragged
            self.dragger.move(event.xdata, event.ydata)
        else: PlotPanel.OnMove(self, event)

    def OnPick(self, event):

        #If a curve is picked, then create and draw an annotation
        x = event.mouseevent.xdata
        y = event.mouseevent.ydata
        thisline = event.artist
        if (not thisline.get_visible()) or self.labeller or self.dragger:
            return

        if event.mouseevent.button == 3:  #Right click
            for state in self.states:
                if thisline in state.curves:  #If a curve is right-clicked, highlight and create tracer
                    flag = True
                    if self.tracer:  #If tracer already exists
                        if thisline == self.tracer.curve:  #if the same curve is clicked
                            self.tracer.stop()
                            self.tracer = None
                            flag = False
                        else:  #if a new curve is clicked
                            self.tracer.setTarget(thisline, state, x, y, reset=True)
                    else:  #create a new tracer if none exists
                        self.tracer = tr.Tracer(self.canvas, self.ax)
                        self.tracer.setTarget(thisline, state, x, y)

                    self.enable_scrolling(not flag)
                elif thisline in state.annotes:  #If an annotation is right-clicked, turn it off
                    state.removeAnnote(thisline)
                    self.redraw()

        elif event.mouseevent.button == 1:   #Left click

            if self.tracer:
                self.tracer.stop()
                self.tracer = None

            for state in self.states:
                if thisline in state.annotes:  #left-click drags annotations
                    self.enable_scrolling(False)
                    thisline.set_visible(False)
                    self.canvas.draw()

                    #Remove the annotation and copy background for blit
                    self.dragger = tr.Mover(self.canvas, self.ax, thisline)
                    self.dragger.start()
                    return
                elif thisline in state.curves: #if a curve is left clicked, then add an annotation
                    i = state.curves.index(thisline)
                    if state.annotes[i] == None:
                        self.enable_scrolling(False)

                        idx = event.ind
                        xd = thisline.get_xdata(orig=True)[idx[0]]
                        yd = thisline.get_ydata(orig=True)[idx[0]]
                        pos = (xd,yd)

                        self.canvas.draw()
                        self.labeller = tr.LineMover(self.canvas, self.ax, self.annoline, pos, state, thisline)
                        self.labeller.start()
                    return
   
    def updateWindow(self):
        PlotPanel.updateWindow(self)
        self.diss.update(self.ax)
    
    def resetAxes(self,  props):
        PlotPanel.resetAxes(self,  props)
        self.diss.update(self.ax)
        
    def clearPlot(self):
        self.ax.cla()

    def OnRelease(self, event):
        if self.labeller:
            tc = (event.xdata, event.ydata)
            state, curve = self.labeller.getData()
            state.create_annote(self.ax, curve, self.labeller.pos, tc)
            self.labeller.stop(tc[0], tc[1])
            self.labeller = None
            self.enable_scrolling(True)
        elif self.dragger:
            self.dragger.stop(event.xdata, event.ydata, show=True)
            self.dragger = None
            self.enable_scrolling(True)


if __name__ == "__main__":
    app = wx.PySimpleApp()

    parser = argparse.ArgumentParser(description='PecPlot Command Line Options')
    parser.add_argument('-i','--infile', help='Input file name',required=False,default='')
    args = parser.parse_args()
    filename = args.infile
    filename = 'test.pec'
    frame = wx.Frame(None, wx.ID_ANY, "")
    panel = TreePlotPanel(frame, filename)
    app.SetTopWindow(frame)
    frame.SetSize((800, 600))
    frame.Show()
    app.MainLoop()
